import os
import csv
import numpy as np
import matplotlib.pylab as plt
from matplotlib.pyplot import plot, ion, show, savefig, cla, figure


# this function load one .cvs (a sequence)
def load_data(dataset, csv_folder='../datasets/NAB-known-anomaly/csv-files/'):
    if dataset == '3dprint':
        data_file = os.path.join(csv_folder, '3dprint.csv')
        anomalies = ['2024-03-16 21:02:55.551', '2024-03-16 21:03:15.068', '2024-03-16 21:04:20.163', '2024-03-16 21:04:31.076', '2024-03-16 21:05:09.843', '2024-03-16 21:05:30.100', '2024-03-16 21:06:13.915', '2024-03-16 21:06:30.347', '2024-03-16 21:06:50.405', '2024-03-16 21:07:01.076']
        t_unit = '0.1 sec'


    t = []
    readings = []
    idx_anomaly = []
    i = 0
    with open(data_file) as csvfile:
        readCSV = csv.reader(csvfile, delimiter=',')
        print("\n--> Anomalies occur at:")
        for row in readCSV:
            if i > 0:
                t.append(i)
                readings.append(float(row[1]))
                for j in range(len(anomalies)):
                    if row[0] == anomalies[j]:
                        idx_anomaly.append(i)
                        print("  timestamp #{}: {}".format(j, row[0]))
            i = i + 1
    t = np.asarray(t)
    readings = np.asarray(readings)
    print("\nOriginal csv file contains {} timestamps.".format(t.shape))
    print("Processed time series contain {} readings.".format(readings.shape))
    print("Anomaly indices are {}".format(idx_anomaly))

    return t, t_unit, readings, idx_anomaly


# This function plots a dataset with the train/test split and known anomalies
# Relies on helper function load_data()

def process_and_save_specified_dataset(dataset, idx_split, y_scale=5, save_file=False):
    t, t_unit, readings, idx_anomaly = load_data(dataset)
    readings = readings.astype(float)

    # split into training and test sets
    training = readings[idx_split[0]:idx_split[1]]
    t_train = t[idx_split[0]:idx_split[1]]

    # normalise by training mean and std
    train_m = np.mean(training)
    train_std = np.std(training)
    print("\nTraining set mean is {}".format(train_m))
    print("Training set std is {}".format(train_std))
    readings_normalised = (readings - train_m) / train_std

    training = readings_normalised[idx_split[0]:idx_split[1]]
    if idx_split[0] == 0:
        test = readings_normalised[idx_split[1]:]
        t_test = t[idx_split[1]:] - idx_split[1]
        idx_anomaly_test = np.asarray(idx_anomaly) - idx_split[1]
    else:
        test = [readings_normalised[:idx_split[0]], readings_normalised[idx_split[1]:]]
        t_test = [t[:idx_split[0]], t[idx_split[1]:] - idx_split[1]]
        idx_anomaly_split = np.squeeze(np.argwhere(np.asarray(idx_anomaly) > idx_split[0]))
        idx_anomaly_test = [np.asarray(idx_anomaly[:idx_anomaly_split[0]]),
                            np.asarray(idx_anomaly[idx_anomaly_split[0]:]) - idx_split[1]]
    print("Anomaly indices in the test set are {}".format(idx_anomaly_test))

    if save_file:
        save_dir = '../datasets/NAB-known-anomaly/'
        np.savez(save_dir + dataset + '.npz', t=t, t_unit=t_unit, readings=readings, idx_anomaly=idx_anomaly,
                 idx_split=idx_split, training=training, test=test, train_m=train_m, train_std=train_std,
                 t_train=t_train, t_test=t_test, idx_anomaly_test=idx_anomaly_test)
        print("\nProcessed time series are saved at {}".format(save_dir + dataset + '.npz'))
    else:
        print("\nProcessed time series are not saved.")

    # plot the whole normalised sequence
    fig, axs = plt.subplots(1, 1, figsize=(18, 4), edgecolor='k')
    fig.subplots_adjust(hspace=.4, wspace=.4)
    # axs = axs.ravel()
    # for i in range(4):
    axs.plot(t, readings_normalised)
    if idx_split[0] == 0:
        axs.plot(idx_split[1] * np.ones(20), np.linspace(-y_scale, y_scale, 20), 'b--')
    else:
        for i in range(2):
            axs.plot(idx_split[i] * np.ones(20), np.linspace(-y_scale, y_scale, 20), 'b--')
    for j in range(len(idx_anomaly)):
        axs.plot(idx_anomaly[j] * np.ones(20), np.linspace(-y_scale, y_scale, 20), 'r--')
    #     axs.plot(data[:,1])
    axs.grid(True)
    axs.set_xlim(0, len(t))
    axs.set_ylim(-y_scale, y_scale)
    axs.set_xlabel("timestamp (every {})".format(t_unit))
    axs.set_ylabel("normalised readings")
    axs.set_title("{} dataset\n(normalised by train mean {:.2f} and std {:.2f})".format(dataset, train_m, train_std))
    axs.legend(('data', 'train test set split', 'anomalies'))
    plt.show()

    return t, readings_normalised

def main():
    process_and_save_specified_dataset('3dprint',[0,5280],y_scale=5,save_file=True)
    # load_data('DB_JIANGYIN_5000','../datasets/NAB-known-anomaly/csv-files/')

if __name__ == '__main__':
    main()